import cv2
import rembg
from tqdm import tqdm
import argparse
import numpy as np

def process_video_with_rembg(args):
    """
    使用rembg处理视频，移除背景并生成一个前景位于纯黑背景上的视频。
    """
    # To use GPU, specify CUDAExecutionProvider
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    print(f"正在准备模型: {args.model} (使用计算后端: {providers})...")
    # 为指定模型创建一个会话，可以提高处理速度
    session = rembg.new_session(model_name=args.model, providers=providers)

    # 准备视频读取
    cap = cv2.VideoCapture(args.input_video)
    if not cap.isOpened():
        print(f"错误: 无法打开视频文件 {args.input_video}")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    if not fps or fps <= 0:
        print("警告: 无法读取有效帧率, 将使用默认值 30 FPS。")
        fps = 30

    print(f"输入视频: {width}x{height}, {fps:.2f} FPS, 共 {total_frames} 帧。")

    # 准备视频写入
    writer = cv2.VideoWriter(
        args.output_file,
        cv2.VideoWriter_fourcc(*'mp4v'),
        fps,
        (width, height)
    )
    if not writer.isOpened():
        print(f"错误: 无法创建视频写入器，请检查路径 '{args.output_file}' 是否有效。")
        cap.release()
        return

    # 逐帧处理
    print("开始处理视频帧...")
    for _ in tqdm(range(total_frames), desc=f"使用 {args.model} 模型处理中"):
        ret, frame = cap.read()
        if not ret:
            break
        
        # 使用 rembg 移除背景, 返回 RGBA 格式
        output_data = rembg.remove(frame, session=session)
        
        # rembg返回的是RGBA，我们需要将其转换为BGR以供OpenCV写入
        # OpenCV处理带alpha通道的图像时，会将透明部分变为黑色，正好符合我们的需求
        # 首先，分离Alpha通道
        alpha = output_data[:, :, 3]
        # 创建一个三通道的BGR图像的可写副本
        bgr_frame = output_data[:, :, :3].copy()
        # 将透明区域设置为黑色
        bgr_frame[alpha == 0] = [0, 0, 0]

        writer.write(bgr_frame)

    # 清理和收尾
    cap.release()
    writer.release()

    print(f"\n处理完成！视频已保存至: {args.output_file}")


if __name__ == "__main__":
    # 直接设置参数
    class Args:
        def __init__(self):
            self.input_video = "C:/Users/Missi/Videos/素材/nina.mp4"  # 修改为你的输入视频路径
            self.output_file = "C:/Users/Missi/Videos/素材/nina_out—1.mp4"  # 修改为你的输出视频路径
            self.model = "isnet-anime"  # 处理动漫推荐 'isnet-anime'，通用场景可选 'u2net'
    设置
    args = Args()
    process_video_with_rembg(args) 